#include "correction_model.h"
#include "../../util.h"
#include "../config.h"

CorrectionModel::CorrectionModel(DistCalc *dist_calc)
{
	this->dist_calc = dist_calc;
}

void CorrectionModel::build(map<string, int> &lm, map<string, int> &tm)
{
	this->lang_model = lm;
	this->trans_model = tm;
	total_freq = 0;
	map<string, int>::iterator itr;
	for(itr = lang_model.begin(); itr != lang_model.end(); ++ itr)
		total_freq += (*itr).second;
	for(itr = trans_model.begin(); itr != trans_model.end(); ++ itr)
	{
		string trans = (*itr).first;
		int freq = (*itr).second;
		int del_trans = trans.find(">");
		string ori = trans.substr(0, del_trans);
		string corr = trans.substr(del_trans + 1, trans.length() - del_trans - 1);
		
		map<string, int>::iterator ori_itr = ori2freq.find(ori);
		if(ori_itr == ori2freq.end())
			ori2freq.insert(pair<string, int>(ori, freq));
		else
			(*ori_itr).second += freq;

		map<string, int>::iterator corr_itr = corr2freq.find(corr);
		if(corr_itr == corr2freq.end())
			corr2freq.insert(pair<string, int>(corr, freq));
		else
			(*corr_itr).second += freq;

		map<string, set<string>>::iterator ori2corr_itr = ori2corr.find(ori);
		if(ori2corr_itr == ori2corr.end())
		{
			set<string> s;
			s.insert(corr);
			ori2corr.insert(pair<string, set<string>>(ori, s));
		}
		else
			(*ori2corr_itr).second.insert(corr);
	}
}

pair<string, double> CorrectionModel::get_prob(vector<string> v, int s, int e)
{
	string cand;
	string ori;
	for(int i = s; i < e; ++ i){
		cand += v[i];
		cand += " ";
	}
	ori = cand;
	cand = util::sort_v(cand);
	//cout << cand << "#\t";
	map<string, set<string>>::iterator p2c_itr = ori2corr.find(cand);
	if(p2c_itr == ori2corr.end()){
		if(corr2freq.find(ori) != corr2freq.end())
			return pair<string, double>(ori, (e-s)*(-100));
		else
			return pair<string, double>("", (e-s)*(-100));
	}
	set<string> corr_s = (*p2c_itr).second;
	set<string>::iterator corr_s_itr;
	double score = NEG_INF, normal_1 = 0, normal_2 = total_freq, t_score, trans_freq, corr_freq;
	normal_1 = (*ori2freq.find(cand)).second;
	string res;
	for(corr_s_itr = corr_s.begin(); corr_s_itr != corr_s.end(); ++ corr_s_itr){
		string corr = *corr_s_itr;	
		string corr_code = util::sort_s(corr);
		string code = cand + ">" + corr;
		trans_freq = (*trans_model.find(code)).second;
		corr_freq = (*lang_model.find(corr_code)).second;
		t_score = pow(trans_freq/normal_1, atof(Config::get("p_alpha").c_str()))
			*pow(corr_freq/normal_2, atof(Config::get("p_beta").c_str()));
		double diff = dist_calc->dist(ori, corr);
		if(abs((double)ori.length() - corr.length()) > atoi(Config::get("hard_term_thr").c_str()))
			continue;
		diff = exp(diff/ori.length());		
		t_score *= pow(diff, atof(Config::get("p_gama").c_str()));
		if(t_score > score){
			score = t_score;
			res = corr;
		}
	}
	return pair<string, double>(res, log(score));
}

string CorrectionModel::correct(string query)
{
	cost[0] = 0;
	map<int, string> align;
	align.insert(pair<int, string>(0, ""));
	vector<string> st = util::tokenize(query, " ");
	for(int i = 1; i <= st.size(); ++ i){
		double score = NEG_INF, t_score, prfx;
		string align_s;
		for(int j = 1; j <= atoi(Config::get("gram").c_str()); ++ j){
			if(j > i)continue;
			int t_prfx = i - j;
			pair<string, double> tmp_p = get_prob(st, t_prfx, i);
			string tmp_s = tmp_p.first;
			double prob = tmp_p.second;
			t_score = cost[t_prfx] + prob;
			//cout << tmp_s << "\t" << prob << "\t" <<t_score << endl;
			if(t_score > score){
				score = t_score;
				prfx = t_prfx;
				align_s = tmp_s;			
			}
		}
		cost[i] = score;
		string s = (*align.find(prfx)).second;
		s += align_s;
		align.insert(pair<int, string>(i, s));
	}
	return (*align.find(st.size())).second;
}

void CorrectionModel::correct_label(map<string, vector<pair<string, int>>> &l2g, map<string, int> &cl)
{
	map<string, vector<pair<string, int>>>::iterator l2g_itr;
	for(l2g_itr = l2g.begin(); l2g_itr != l2g.end(); ++ l2g_itr)
	{
		string correction = correct((*l2g_itr).first);
		cl.insert(pair<string, int>(correction, 1));
	}
}